from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage

SPARKAI_URL = 'wss://spark-api.xf-yun.com/v1.1/chat'
SPARKAI_APP_ID = 'ae6000ef'
SPARKAI_API_SECRET = 'OWQ5MjI3YjI3OTIyNjczNDg3YjAxZDFi'
SPARKAI_API_KEY = 'e263794f4d8fcf6b22f6f6f3f7de7390'
SPARKAI_DOMAIN = 'lite'


if __name__ == '__main__':
    spark = ChatSparkLLM(
        spark_api_url=SPARKAI_URL,
        spark_app_id=SPARKAI_APP_ID,
        spark_api_key=SPARKAI_API_KEY,
        spark_api_secret=SPARKAI_API_SECRET,
        spark_llm_domain=SPARKAI_DOMAIN,
        streaming=False,
    )
    messages = [ChatMessage(
         role="user",
         content='你好呀'
    )]
    handler = ChunkPrintHandler() 
    a = spark.generate([messages], callbacks=[handler])
    print(a)
